import os
import scipy
import numpy as np
import tensorflow as tf
import logging
from config import cfg
import app_logger 

logger = logging.getLogger("utils") 
app_logger.init(cfg.log_lvl)
def load_mnist(path, is_training):
    fd = open(os.path.join(cfg.dataset, 'train-images-idx3-ubyte'), "rb")
    loaded = np.fromfile(file=fd, dtype=np.uint8)
    trX = loaded[16:].reshape((60000, 28, 28, 1)).astype(np.float32)

    fd = open(os.path.join(cfg.dataset, 'train-labels-idx1-ubyte'), "rb")
    loaded = np.fromfile(file=fd, dtype=np.uint8)
    trY = loaded[8:].reshape((60000)).astype(np.int32)

    fd = open(os.path.join(cfg.dataset, 't10k-images-idx3-ubyte'), "rb")
    loaded = np.fromfile(file=fd, dtype=np.uint8)
    teX = loaded[16:].reshape((10000, 28, 28, 1)).astype(np.float32)

    fd = open(os.path.join(cfg.dataset, 't10k-labels-idx1-ubyte'), "rb")
    loaded = np.fromfile(file=fd, dtype=np.uint8)
    teY = loaded[8:].reshape((10000)).astype(np.int32)

    # normalization 
    if is_training:
        return trX/255., trY
    else:
        return teX/255., teY

class MNISTSamples():
    def __init__(self, is_train = True, shuffle = True, batch_size = 32, shift_img = 0):
        x, y = load_mnist(cfg.dataset, is_train)
        if shift_img > 0:
            x_new = [None]*4
            y_new = [None]*4
            for i in range(4):
                xs = self.__shift_by(shift_img, i, x)
                x_new[i] = xs
                y_new[i] = y
            x = np.concatenate([x] + x_new, axis=0)
            y = np.concatenate([y] + y_new, axis=0)
        self.x, self.y = x, y
        self.batch_size = batch_size 
        self.shuffle = shuffle 
        self.sample_count = self.x.shape[0]
        if shuffle:
            idx = np.arange(self.sample_count)
            np.random.shuffle(idx)
            self.x = self.x[idx]
            self.y = self.y[idx]
        append_num = batch_size - self.sample_count % batch_size 
        self.x = np.concatenate((self.x, self.x[:append_num]))
        self.y = np.concatenate((self.y, self.y[:append_num]))
        self.sample_count += append_num
        self.__len = self.sample_count//batch_size
        self.__shuffle_index = np.arange(self.__len)

    def __shift_by(self, shift_size, direction, data_set):
        n = len(data_set)
        x = [None]*n
        direction %= 4
        for i in range(n):
            xs = data_set[i]
            if direction == 0:
                # top
                xs = np.concatenate((xs, np.zeros(
                        (shift_size, xs.shape[1], xs.shape[2]), 
                        dtype = np.float32)),
                     axis = 0)
                x[i] = xs[shift_size:, :, :]
            if direction == 1:
                # left
                xs = np.concatenate((xs, np.zeros(
                        (xs.shape[0], shift_size, xs.shape[2]), 
                        dtype = np.float32)),
                     axis = 1)
                x[i] = xs[:, shift_size:, :]
            if direction == 2:
                # down
                xs = np.concatenate((np.zeros(
                        (shift_size, xs.shape[1], xs.shape[2]), 
                        dtype = np.float32), xs),
                     axis = 0)
                x[i] = xs[:-shift_size, :, :]
            if direction == 3:
                # right
                xs = np.concatenate((np.zeros(
                        (xs.shape[0], shift_size, xs.shape[2]), 
                        dtype = np.float32), xs),
                     axis = 1)
                x[i] = xs[:, :-shift_size, :]
        return np.array(x)

    def __len__(self):
        return self.__len
    
    def __getitem__(self, index):
        index = index % self.__len
        if self.shuffle and index == 0:
            np.random.shuffle(self.__shuffle_index)
        indx = self.__shuffle_index[index]
        start = indx * self.batch_size
        end = start + self.batch_size
        img = self.x[start : end]
        new_img = [None]*self.batch_size
        for i in range(self.batch_size):
            new_img[i] = shift_img_random(img[i], 2)
        return new_img, self.y[start : end]
   
def shift_img_random(img, max_dist):
    x_sh = np.random.randint(-max_dist, max_dist+1)
    y_sh = np.random.randint(-max_dist, max_dist+1)
    height, width, ch = img.shape[0], img.shape[1], img.shape[2]
    img = np.copy(img)
    if x_sh > 0:
        img = img[:, x_sh:]
        img = np.concatenate((img, np.zeros((height, x_sh, ch), dtype=np.float32)), axis=1)
    elif x_sh < 0:
        img = img[:, :x_sh]
        img = np.concatenate((np.zeros((height, -x_sh, ch), dtype=np.float32), img), axis=1)

    if y_sh > 0:
        img = img[y_sh:, :]
        img = np.concatenate((img, np.zeros((y_sh, width, ch), dtype=np.float32)), axis=0)
    elif y_sh < 0:
        img = img[:y_sh, :]
        img = np.concatenate((np.zeros((-y_sh, width, ch), dtype=np.float32), img), axis=0)

    return img

def save_images(imgs, size, path):
    '''
    Args:
        imgs: [batch_size, image_height, image_width]
        size: a list with tow int elements, [image_height, image_width]
        path: the path to save images
    '''
    imgs = (imgs + 1.) / 2  # inverse_transform
    return(scipy.misc.imsave(path, mergeImgs(imgs, size)))


def mergeImgs(images, size):
    h, w = images.shape[1], images.shape[2]
    imgs = np.zeros((h * size[0], w * size[1], 3))
    for idx, image in enumerate(images):
        i = idx % size[1]
        j = idx // size[1]
        imgs[j * h:j * h + h, i * w:i * w + w, :] = image

    return imgs


if __name__ == '__main__':
    from matplotlib import pyplot as plt
    test_sample =  MNISTSamples(is_train = False, shuffle = False, batch_size = cfg.batch_size)
    imgs, label = test_sample[0]
    img = imgs[2]
    _, (ax0, ax1) = plt.subplots(1,2)
    ax1.imshow(np.squeeze(shift_img_random(img, 2)))
    ax0.imshow(np.squeeze(img))
    plt.show()

